import isaacgym
import torch
from asyncio import FastChildWatcher
from code import interact
from collections import deque
import os
import time
from tkinter.tix import Tree
import numpy as np
from isaacgym import gymutil
from isaacgym import gymapi
#from isaacgym import gymtorch
from math import sqrt
import math
from sympy import false
import cv2
from draw import *
from pcgworker.PCGWorker import *

from wfc_vecenv_stable_baselines import PCGVecEnvStableBaselines

from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.results_plotter import load_results, ts2xy
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common import results_plotter
from stable_baselines3.common.results_plotter import load_results, ts2xy, plot_results
from stable_baselines3.common.noise import NormalActionNoise
from stable_baselines3.common.callbacks import BaseCallback

import ray
import torch

import asyncio

from multiprocessing import Process, Queue, Pipe, Manager, Value, Array, Lock, Event, Pool

def evaluate_model(model_path, seed_id, seed_path, device_id, return_queue = None):
    
    os.environ['CUDA_VISIBLE_DEVICES']='0,1,2,3,4,5,6,7'
    LOGDIR = "./training_logs/"

    print(seed_path)

    m_env = PCGVecEnvStableBaselines(headless_ = True,seed_pth = seed_path, compute_device_id = device_id, graphics_device_id = device_id)
    m_env.reset()

    # load model
    model = PPO.load(LOGDIR+model_path, device = 'cuda:'+str(device_id))

    # evaluate
    eval = m_env.evaluate_sb3(model = model, num_episodes = 100)
    print("eval: ", eval)

    m_env.close()

    return_queue.put((
        seed_id,
        eval
    ))


if __name__ == "__main__":

    LOGDIR = "./training_logs"

    manager = multiprocessing.Manager()
    return_q = manager.Queue()

    results = []

    winners_list = [
        "gen_0_dec_3.json",
        "gen_1_dec_4.json",
        "gen_2_dec_5.json",
        "gen_3_dec_5.json",
        "gen_4_dec_3.json",
        "gen_5_dec_4.json",
        "gen_6_dec_7.json",
        "gen_7_dec_7.json",
        "gen_8_dec_5.json",
        "gen_9_dec_1.json",
        "gen_10_dec_1.json",
        "gen_11_dec_5.json",
        "gen_12_dec_1.json",
        "gen_13_dec_2.json",
        "gen_14_dec_2.json",
        "gen_15_dec_4.json",
        "gen_16_dec_0.json",
        "gen_17_dec_3.json",
        "gen_18_dec_5.json",
        "gen_19_dec_0.json",
        "gen_20_dec_1.json",
        "gen_21_dec_4.json",
        "gen_22_dec_0.json",
        "gen_23_dec_4.json",
        "gen_24_dec_5.json",
        "gen_25_dec_3.json",
        "gen_26_dec_4.json",
        "gen_27_dec_7.json",
        "gen_28_dec_1.json",
        "gen_29_dec_2.json",
        "gen_30_dec_1.json",
        "gen_31_dec_2.json",
        "gen_32_dec_6.json",
        "gen_33_dec_6.json",
        "gen_34_dec_0.json",
        "gen_35_dec_0.json",
        "gen_36_dec_3.json",
        "gen_37_dec_4.json",
        "gen_38_dec_0.json",
        "gen_39_dec_1.json",
        "gen_40_dec_7.json",
    ]

    for m in range(0, 41):

        ev_model = "trial2/model_"+str(m)+".zip"

        p0 = Process(target=evaluate_model, args=(ev_model,0,"./trial2/"+winners_list[0],0,return_q,))
        p1 = Process(target=evaluate_model, args=(ev_model,1,"./trial2/"+winners_list[1],1,return_q,))
        p2 = Process(target=evaluate_model, args=(ev_model,2,"./trial2/"+winners_list[2],2,return_q,))
        p3 = Process(target=evaluate_model, args=(ev_model,3,"./trial2/"+winners_list[3],3,return_q,))
        p4 = Process(target=evaluate_model, args=(ev_model,4,"./trial2/"+winners_list[4],4,return_q,))
        p5 = Process(target=evaluate_model, args=(ev_model,5,"./trial2/"+winners_list[5],5,return_q,))
        p6 = Process(target=evaluate_model, args=(ev_model,6,"./trial2/"+winners_list[6],6,return_q,))
        p7 = Process(target=evaluate_model, args=(ev_model,7,"./trial2/"+winners_list[7],7,return_q,))

        p0.start()
        p1.start()
        p2.start()
        p3.start()
        p4.start()
        p5.start()
        p6.start()
        p7.start()

        p0.join()
        p1.join()
        p2.join()
        p3.join()
        p4.join()
        p5.join()
        p6.join()
        p7.join()

        p0.kill()
        p1.kill()
        p2.kill()
        p3.kill()
        p4.kill()
        p5.kill()
        p6.kill()
        p7.kill()

        output_1 = []
        for i in range(return_q.qsize()):
            results = copy.copy(return_q.get())
            return_q.task_done()
            output_1.append(results)

        # print(output_1)

        p0 = Process(target=evaluate_model, args=(ev_model,8,"./trial2/"+winners_list[8],0,return_q,))
        p1 = Process(target=evaluate_model, args=(ev_model,9,"./trial2/"+winners_list[9],1,return_q,))
        p2 = Process(target=evaluate_model, args=(ev_model,10,"./trial2/"+winners_list[10],2,return_q,))
        p3 = Process(target=evaluate_model, args=(ev_model,11,"./trial2/"+winners_list[11],3,return_q,))
        p4 = Process(target=evaluate_model, args=(ev_model,12,"./trial2/"+winners_list[12],4,return_q,))
        p5 = Process(target=evaluate_model, args=(ev_model,13,"./trial2/"+winners_list[13],5,return_q,))
        p6 = Process(target=evaluate_model, args=(ev_model,14,"./trial2/"+winners_list[14],6,return_q,))
        p7 = Process(target=evaluate_model, args=(ev_model,15,"./trial2/"+winners_list[15],7,return_q,))

        p0.start()
        p1.start()
        p2.start()
        p3.start()
        p4.start()
        p5.start()
        p6.start()
        p7.start()

        p0.join()
        p1.join()
        p2.join()
        p3.join()
        p4.join()
        p5.join()
        p6.join()
        p7.join()

        p0.kill()
        p1.kill()
        p2.kill()
        p3.kill()
        p4.kill()
        p5.kill()
        p6.kill()
        p7.kill()

        output_2 = []
        for i in range(return_q.qsize()):
            results = copy.copy(return_q.get())
            return_q.task_done()
            output_2.append(results)

        # print(output_2)

        p0 = Process(target=evaluate_model, args=(ev_model,16,"./trial2/"+winners_list[16],0,return_q,))
        p1 = Process(target=evaluate_model, args=(ev_model,17,"./trial2/"+winners_list[17],1,return_q,))
        p2 = Process(target=evaluate_model, args=(ev_model,18,"./trial2/"+winners_list[18],2,return_q,))
        p3 = Process(target=evaluate_model, args=(ev_model,19,"./trial2/"+winners_list[19],3,return_q,))
        p4 = Process(target=evaluate_model, args=(ev_model,20,"./trial2/"+winners_list[20],4,return_q,))
        p5 = Process(target=evaluate_model, args=(ev_model,21,"./trial2/"+winners_list[21],5,return_q,))
        p6 = Process(target=evaluate_model, args=(ev_model,22,"./trial2/"+winners_list[22],6,return_q,))
        p7 = Process(target=evaluate_model, args=(ev_model,23,"./trial2/"+winners_list[23],7,return_q,))

        p0.start()
        p1.start()
        p2.start()
        p3.start()
        p4.start()
        p5.start()
        p6.start()
        p7.start()

        p0.join()
        p1.join()
        p2.join()
        p3.join()
        p4.join()
        p5.join()
        p6.join()
        p7.join()

        p0.kill()
        p1.kill()
        p2.kill()
        p3.kill()
        p4.kill()
        p5.kill()
        p6.kill()
        p7.kill()

        output_3 = []
        for i in range(return_q.qsize()):
            results = copy.copy(return_q.get())
            return_q.task_done()
            output_3.append(results)

        # print(output_3)

        p0 = Process(target=evaluate_model, args=(ev_model,24,"./trial2/"+winners_list[24],0,return_q,))
        p1 = Process(target=evaluate_model, args=(ev_model,25,"./trial2/"+winners_list[25],1,return_q,))
        p2 = Process(target=evaluate_model, args=(ev_model,26,"./trial2/"+winners_list[26],2,return_q,))
        p3 = Process(target=evaluate_model, args=(ev_model,27,"./trial2/"+winners_list[27],3,return_q,))
        p4 = Process(target=evaluate_model, args=(ev_model,28,"./trial2/"+winners_list[28],4,return_q,))
        p5 = Process(target=evaluate_model, args=(ev_model,29,"./trial2/"+winners_list[29],5,return_q,))
        p6 = Process(target=evaluate_model, args=(ev_model,30,"./trial2/"+winners_list[30],6,return_q,))
        p7 = Process(target=evaluate_model, args=(ev_model,31,"./trial2/"+winners_list[31],7,return_q,))

        p0.start()
        p1.start()
        p2.start()
        p3.start()
        p4.start()
        p5.start()
        p6.start()
        p7.start()

        p0.join()
        p1.join()
        p2.join()
        p3.join()
        p4.join()
        p5.join()
        p6.join()
        p7.join()

        p0.kill()
        p1.kill()
        p2.kill()
        p3.kill()
        p4.kill()
        p5.kill()
        p6.kill()
        p7.kill()

        output_4 = []
        for i in range(return_q.qsize()):
            results = copy.copy(return_q.get())
            return_q.task_done()
            output_4.append(results)

        # print(output_4)

        p0 = Process(target=evaluate_model, args=(ev_model,32,"./trial2/"+winners_list[32],0,return_q,))
        p1 = Process(target=evaluate_model, args=(ev_model,33,"./trial2/"+winners_list[33],1,return_q,))
        p2 = Process(target=evaluate_model, args=(ev_model,34,"./trial2/"+winners_list[34],2,return_q,))
        p3 = Process(target=evaluate_model, args=(ev_model,35,"./trial2/"+winners_list[35],3,return_q,))
        p4 = Process(target=evaluate_model, args=(ev_model,36,"./trial2/"+winners_list[36],4,return_q,))
        p5 = Process(target=evaluate_model, args=(ev_model,37,"./trial2/"+winners_list[37],5,return_q,))
        p6 = Process(target=evaluate_model, args=(ev_model,38,"./trial2/"+winners_list[38],6,return_q,))
        p7 = Process(target=evaluate_model, args=(ev_model,39,"./trial2/"+winners_list[39],7,return_q,))
            
        p0.start()
        p1.start()
        p2.start()
        p3.start()
        p4.start()
        p5.start()
        p6.start()
        p7.start()

        p0.join()
        p1.join()
        p2.join()
        p3.join()
        p4.join()
        p5.join()
        p6.join()
        p7.join()

        p0.kill()
        p1.kill()
        p2.kill()
        p3.kill()
        p4.kill()
        p5.kill()
        p6.kill()
        p7.kill()

        output_5 = []
        for i in range(return_q.qsize()):
            results = copy.copy(return_q.get())
            return_q.task_done()
            output_5.append(results)

        # print(output_5)

        p0 = Process(target=evaluate_model, args=(ev_model,40,"./trial2/"+winners_list[40],0,return_q,))
        p0.start()
        p0.join()
        p0.kill()
        
        output_6 = []
        for i in range(return_q.qsize()):
            results = copy.copy(return_q.get())
            return_q.task_done()
            output_6.append(results)

        # concatenate all the results
        output = output_1 + output_2 + output_3 + output_4 + output_5 + output_6
        # sort the results by the seed_id
        output.sort(key=lambda x: x[0])

        # print the results
        print("-------------------------"+'eval_model_'+str(m)+'_results-------------------------')
        for i in range(len(output)):
            print(output[i])

        # save the results to a json file
        with open('eval_model_'+str(m)+'.json' , 'w') as fp:
            json.dump(output, fp)